package com.oliver.proxy.midjourney.application.service;

import com.alibaba.fastjson2.JSONObject;
import com.oliver.proxy.midjourney.domain.entity.Prod;
import com.oliver.proxy.midjourney.domain.entity.TaskList;
import com.oliver.proxy.midjourney.infrastructure.mapper.ProdMapper;
import com.oliver.proxy.midjourney.infrastructure.mapper.TaskListMapper;
import io.github.yuanbaobaoo.dify.client.DifyClientBuilder;
import io.github.yuanbaobaoo.dify.client.IDifyWorkFlowClient;
import io.github.yuanbaobaoo.dify.client.params.ParamMessage;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.annotation.PostConstruct;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.stream.Collectors;

@Slf4j
@Service
public class PromptProcessingService {
    private final ProdMapper prodMapper;
    private final TaskListMapper taskListMapper;
    private final ExecutorService executor;
    private final AtomicBoolean isProcessing = new AtomicBoolean(false);
    private static final int THREAD_COUNT = 4;
    private final List<IDifyWorkFlowClient> difyClients = new ArrayList<>();

    @Autowired
    public PromptProcessingService(ProdMapper prodMapper, TaskListMapper taskListMapper) {
        this.prodMapper = prodMapper;
        this.taskListMapper = taskListMapper;
        
        // Create multiple Dify client instances, one per thread
        for (int i = 0; i < THREAD_COUNT; i++) {
            difyClients.add(DifyClientBuilder.create()
                    .apiKey("app-LhVH6zMehG7JLAx2t1unCOFE")
                    .baseUrl("http://localhost/v1")
                    .buildWorkFlow());
        }
        
        this.executor = Executors.newFixedThreadPool(THREAD_COUNT);
    }

    @PostConstruct
    public void init() {
        CompletableFuture.runAsync(this::processPromptsInBackground);
    }

    @Scheduled(fixedRate = 6000)
    public void checkAndRestartProcessing() {
        if (!isProcessing.get()) {
            log.info("Restarting prompt processing...");
            CompletableFuture.runAsync(this::processPromptsInBackground);
        }
    }

    private void processPromptsInBackground() {
        try {
            if (!isProcessing.compareAndSet(false, true)) {
                return;
            }
            
            log.info("Starting prompt processing in background thread");
            
            while (true) {
                // Get all task lists with type 1 and their associated prods with empty prompts
                List<Map<String, Object>> taskData = taskListMapper.selectTasksWithEmptyPrompts();
                
                if (taskData.isEmpty()) {
                    Thread.sleep(3000);
                    continue;
                }

                log.info("Found {} tasks to process", taskData.size());
                
                // Process tasks in parallel
                processTasksInParallel(taskData);
            }
        } catch (Exception e) {
            log.error("Error in prompt background processing", e);
        } finally {
            isProcessing.set(false);
        }
    }

    private void processTasksInParallel(List<Map<String, Object>> taskData) {
        // Group tasks by task_id
        Map<Long, List<Map<String, Object>>> taskGroups = taskData.stream()
                .collect(Collectors.groupingBy(data -> Long.parseLong(data.get("task_id").toString())));

        List<CompletableFuture<Void>> futures = new ArrayList<>();

        taskGroups.forEach((taskId, prods) -> {
            int switchNum = (Integer.parseInt(prods.get(0).get("task_switch_num").toString()))*4;
            
            // Split prods into batches based on task_switch_num
            List<List<Map<String, Object>>> batches = new ArrayList<>();
            for (int i = 0; i < prods.size(); i += switchNum) {
                batches.add(prods.subList(i, Math.min(i + switchNum, prods.size())));
            }

            // Process each batch
            for (List<Map<String, Object>> batch : batches) {
                CompletableFuture<Void> future = CompletableFuture.runAsync(() -> 
                    processBatch(batch), executor);
                futures.add(future);
            }
        });

        CompletableFuture.allOf(futures.toArray(new CompletableFuture[0])).join();
    }

    private void processBatch(List<Map<String, Object>> batch) {
        try {
            // Get a Dify client from the pool
            IDifyWorkFlowClient client = difyClients.get(ThreadLocalRandom.current().nextInt(THREAD_COUNT));

            // Prepare batch information for Dify
            String taskPrompt = batch.get(0).get("task_prompt").toString();
            String imgSize = batch.get(0).get("img_size").toString();
            List<Long> prodIds = batch.stream()
                    .map(prod -> Long.parseLong(prod.get("data_id").toString()))
                    .collect(Collectors.toList());

            // Create message for Dify
            ParamMessage message = ParamMessage.builder()
                    .user("x@oliverding.cn")
                    .inputs(new HashMap<>() {{
                        put("prompt", taskPrompt);
                        put("size", imgSize);
                    }})
                    .build();

            // Call Dify API
            JSONObject result = client.runBlocking(message);

            if (result != null && result.getJSONObject("data") != null) {
                JSONObject data = result.getJSONObject("data");
                if ("succeeded".equals(data.getString("status"))) {
                    String prompt = data.getJSONObject("outputs").getString("text");
                    updateProdsPrompt(prodIds, prompt);
                    log.info("Successfully updated prompt for batch with {} prods", batch.size());
                }
            }

            Thread.sleep(1000);
        } catch (Exception e) {
            log.error("Error processing batch: ", e);
        }
    }

    @Transactional
    public void updateProdsPrompt(List<Long> prodIds, String prompt) {
        prodMapper.updateBatchPrompt(prodIds, prompt);
    }
} 